# -*- coding: utf-8 -*-
# Scikit-learn-2.py
"""
#### 题目2: 随机森林分类鸢尾花（分类问题）
**任务描述**: 
使用鸢尾花数据集(Iris)训练一个随机森林分类器。  
**步骤**:  
1. 加载鸢尾花数据集。  
2. 划分训练集和测试集(80%训练, 20%测试)。  
3. 训练随机森林分类器(设置n_estimators=100)。  
4. 在测试集上进行预测。  
5. 使用混淆矩阵和分类报告(包括精确率、召回率、F1分数)评估模型。  
6. 绘制ROC曲线(注意: 这是一个多分类问题, ROC曲线需要为每个类别分别绘制)。  
"""

from matplotlib import pyplot as plt
from sklearn.datasets import load_iris                  # 鸢尾花数据集
from sklearn.ensemble import RandomForestClassifier     # 随机森林分类器
from sklearn.metrics import ConfusionMatrixDisplay      # 混淆矩阵
from sklearn.metrics import classification_report       # 分类报告
from sklearn.metrics import roc_curve, auc              # ROC曲线
from sklearn.preprocessing import label_binarize
from ScikitLearnBase import ScikitLearnExperiment      # 导入实验基类

# 创建实验实例
experiment = ScikitLearnExperiment("随机森林分类鸢尾花")

# 打印实验标题
experiment.print_header()

# 1. 加载鸢尾花数据集
print("1. 加载鸢尾花数据集...")
iris = experiment.load_dataset(load_iris)

# 创建中英文映射字典
class_name_mapping = {
    'setosa': '山鸢尾',
    'versicolor': '变色鸢尾',
    'virginica': '维吉尼亚鸢尾'
}

feature_name_mapping = {
    'sepal length (cm)': '萼片长度（厘米）',
    'sepal width (cm)': '萼片宽度（厘米）',
    'petal length (cm)': '花瓣长度（厘米）',
    'petal width (cm)': '花瓣宽度（厘米）'
}

# 创建中文名称列表
class_names_zh = experiment.create_target_name_mapping(class_name_mapping)
feature_names_zh = experiment.create_feature_name_mapping(feature_name_mapping)

print(f"类别名称(中文): {', '.join(class_names_zh)}")
print(f"特征名称(中文): {', '.join(feature_names_zh)}")

# 2. 划分训练集和测试集（80%训练，20%测试）
experiment.split_data(test_size=0.2, random_state=0)

# 3. 训练随机森林分类器，设置n_estimators=100
print("3. 训练随机森林分类器中...")
clf = RandomForestClassifier(n_estimators=100, max_depth=2, random_state=0)
clf.fit(experiment.X_train, experiment.y_train)
print("训练完成！")

# 4. 在测试集上进行预测
y_pred = clf.predict(experiment.X_test)
# 获取预测概率
y_pred_proba = clf.predict_proba(experiment.X_test)
print("预测结果:", y_pred)

# 5. 评估模型
accuracy = experiment.evaluate_model(clf, experiment.X_test, experiment.y_test, "随机森林分类器")

# 5.1 使用混淆矩阵评估模型
print("\n混淆矩阵:")
fig, ax = plt.subplots(figsize=(8, 6))
disp = ConfusionMatrixDisplay.from_estimator(
    clf, experiment.X_test, experiment.y_test,
    display_labels=class_names_zh,   # 使用中文类别名称
    cmap=plt.cm.Blues,
    ax=ax
)
disp.ax_.set_title("鸢尾花分类混淆矩阵")
plt.tight_layout()

# 保存混淆矩阵图片
experiment.save_plot("iris_confusion_matrix.png")
plt.show()

# 5.2 使用分类报告评估模型
print("\n分类报告:")
print(classification_report(experiment.y_test, y_pred, target_names=class_names_zh))

# 6. 绘制ROC曲线（多分类）
# 将标签二值化
y_test_bin = label_binarize(experiment.y_test, classes=[0, 1, 2])
n_classes = y_test_bin.shape[1]

# 计算每个类的ROC曲线和AUC
fpr = dict()
tpr = dict()
roc_auc = dict()
for i in range(n_classes):
    fpr[i], tpr[i], _ = roc_curve(y_test_bin[:, i], y_pred_proba[:, i])
    roc_auc[i] = auc(fpr[i], tpr[i])

# 绘制所有ROC曲线
plt.figure(figsize=(10, 8))
colors = ['blue', 'red', 'green']
for i, color in zip(range(n_classes), colors):
    plt.plot(fpr[i], tpr[i], color=color, lw=2,
             label='ROC curve of class {0} (AUC = {1:0.2f})'
             ''.format(class_names_zh[i], roc_auc[i]))   # 使用中文类别名称

plt.plot([0, 1], [0, 1], 'k--', lw=2)
plt.xlim([0.0, 1.0])
plt.ylim([0.0, 1.05])
plt.xlabel('假正率 (False Positive Rate)')
plt.ylabel('真正率 (True Positive Rate)')
plt.title('多类别ROC曲线')
plt.legend(loc="lower right")
plt.grid(alpha=0.3)
plt.tight_layout()

# 保存ROC曲线图片
experiment.save_plot("iris_roc_curve.png")
plt.show()

# 打印实验结束信息
experiment.print_footer()